import os
from PIL import Image
from torch.utils.data import Dataset


class CUB(Dataset):

    def __init__(self, data_path, train, transform=None):
        self.root = data_path
        self.transform = transform
        self.id2path = {}
        with open(os.path.join(self.root, 'images.txt')) as f:
            for line in f:
                image_id, path = line.split()
                self.id2path[image_id] = path

        self.id2class = {}
        with open(os.path.join(self.root, 'image_class_labels.txt')) as f:
            for line in f:
                image_id, class_id = line.split()
                self.id2class[image_id] = int(class_id) - 1

        self.data_id = []
        if train:
            with open(os.path.join(self.root, 'train_test_split.txt')) as f:
                for line in f:
                    image_id, is_train = line.split()
                    if int(is_train):
                        self.data_id.append(image_id)
        else:
            with open(os.path.join(self.root, 'train_test_split.txt')) as f:
                for line in f:
                    image_id, is_train = line.split()
                    if not int(is_train):
                        self.data_id.append(image_id)

    def __len__(self):
        return len(self.data_id)

    def __getitem__(self, index):
        image_id = self.data_id[index]
        path = self.id2path[image_id]
        image = Image.open(os.path.join(self.root, 'images', path)).convert('RGB')
        label = self.id2class[image_id]
        if self.transform:
            image = self.transform(image)
        return image, label


class Dogs(Dataset):

    def __init__(self, data_path, train, transform=None):
        self.root = data_path
        self.transform = transform
        self.path_label = {}
        with open(os.path.join(self.root, 'labels.txt')) as f:
            for line in f:
                path, label = line.split()
                self.path_label[path] = int(label) - 1

        self.data_id = []
        if train:
            with open(os.path.join(self.root, 'train_data.txt')) as f:
                self.data_id = f.read().splitlines()

        else:
            with open(os.path.join(self.root, 'test_data.txt')) as f:
                self.data_id = f.read().splitlines()

    def __len__(self):
        return len(self.data_id)

    def __getitem__(self, index):
        path = self.data_id[index]
        label = self.path_label[path]
        image = Image.open(os.path.join(self.root, 'Images', path)).convert('RGB')

        if self.transform:
            image = self.transform(image)
        return image, label


class Aircrafts(Dataset):
    def __init__(self, data_path, train, transform=None):
        self.root = data_path
        self.transform = transform
        self.index = None

        with open(os.path.join(data_path, 'variants.txt'), 'r') as file:
            self.category = {line.strip(): index for index, line in enumerate(file.readlines())}

        if train:
            with open(os.path.join(data_path, 'images_variant_trainval.txt'), "r") as f:
                self.data = [line.strip().split(" ", maxsplit=1) for line in f]

        else:
            with open(os.path.join(data_path, 'images_variant_test.txt'), "r") as f:
                self.data = [line.strip().split(" ", maxsplit=1) for line in f]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        path = os.path.join(self.root, 'images', self.data[idx][0] + '.jpg')
        image = Image.open(path).convert('RGB')
        label = self.category[self.data[idx][1]]
        if self.transform:
            image = self.transform(image)
        return image, label

# if __name__ == '__main__':
#     data = Aircrafts(r'E:\desktop\airs',train=False)
#     x, y = data[100]
#     print(y)
#     x.show()
